#
# Created: 2025-08-06

import plotly.graph_objects as go
from plotly.subplots import make_subplots
import pandas as pd
from matplotlib.colors import LinearSegmentedColormap, to_hex, to_rgb

# 读取数据
df = pd.read_csv('sales_success.csv')

# 层级定义
levels = ['salesperson', 'county', 'region']
color_columns = ['sales', 'calls']
value_column = 'calls'

# 构建多层级数据结构
def build_hierarchical_dataframe(df, levels, value_column, color_columns=None):
    df_all_trees = []
    region_list, county_list, salesperson_list = [], [], []

    for i, level in enumerate(levels):
        df_tree = pd.DataFrame(columns=['id', 'parent', 'value', 'color'])
        dfg = df.groupby(levels[i:]).sum().reset_index()

        df_tree['id'] = dfg[level].copy()
        if i < len(levels) - 1:
            df_tree['parent'] = dfg[levels[i+1]].copy()
        else:
            df_tree['parent'] = 'total'

        df_tree['value'] = dfg[value_column]
        df_tree['color'] = dfg[color_columns[0]] / dfg[color_columns[1]]
        df_all_trees.append(df_tree)

        if level == 'region':
            region_list = df_tree['id'].tolist()
        elif level == 'county':
            county_list = df_tree['id'].tolist()
        elif level == 'salesperson':
            salesperson_list = df_tree['id'].tolist()

    # 添加根节点
    total = pd.Series(dict(
        id='total',
        parent='',
        value=df[value_column].sum(),
        color=df[color_columns[0]].sum() / df[color_columns[1]].sum()
    ))
    df_all_trees.append(pd.DataFrame(total).T)

    df_all = pd.concat(df_all_trees, ignore_index=True)
    return df_all, region_list, county_list, salesperson_list

# 构建结构
df_all_trees, region_list, county_list, salesperson_list = build_hierarchical_dataframe(
    df, levels, value_column, color_columns
)

# 原始基础色
base_colors = ["#FDDED7", "#F5BE8F", "#C1E0DB", "#CCD376"]

# 提亮函数
def lighten_color(hex_color, factor=0.35):
    rgb = to_rgb(hex_color)
    light_rgb = [1 - (1 - c) * (1 - factor) for c in rgb]
    return to_hex(light_rgb)

# 提亮 base colors
lighter_base_colors = [lighten_color(c, factor=0.35) for c in base_colors]
lighter_cmap = LinearSegmentedColormap.from_list("lighter", lighter_base_colors)

# 为每一层生成颜色
county_colors = [to_hex(lighter_cmap(i / (len(county_list) - 1))) for i in range(len(county_list))]
salesperson_colors = [to_hex(lighter_cmap(i / (len(salesperson_list) - 1))) for i in range(len(salesperson_list))]
region_colors = lighter_base_colors

# 构建颜色映射 color_map
color_map = {"total": "#FFFFFF"}

for i, name in enumerate(region_list):
    color_map[name] = region_colors[i % len(region_colors)]

for i, name in enumerate(county_list):
    color_map[name] = county_colors[i]

for i, name in enumerate(salesperson_list):
    color_map[name] = salesperson_colors[i]

# 映射颜色列表
color_list = df_all_trees['id'].map(color_map)

# 创建双图结构
fig = make_subplots(1, 2, specs=[[{"type": "domain"}, {"type": "domain"}]])

# 图1：完整层级
fig.add_trace(go.Sunburst(
    labels=df_all_trees['id'],
    parents=df_all_trees['parent'],
    values=df_all_trees['value'],
    branchvalues='total',
    marker=dict(colors=color_list),
    hovertemplate='<b>%{label}</b><br>Sales: %{value}<br>Success rate: %{color:.2f}'
), 1, 1)

# 图2：限制2层
fig.add_trace(go.Sunburst(
    labels=df_all_trees['id'],
    parents=df_all_trees['parent'],
    values=df_all_trees['value'],
    branchvalues='total',
    marker=dict(colors=color_list),
    hovertemplate='<b>%{label}</b><br>Sales: %{value}<br>Success rate: %{color:.2f}',
    maxdepth=2
), 1, 2)

# 设置图表边距
fig.update_layout(margin=dict(t=10, b=10, r=10, l=10))
fig.write_html("sunburst_chart.html")
fig.show()

